import torch
from skorch.layers import *
from skorch.vision import *
from skorch.model import create_cnn

def regressor_cnn(n_inputs=3, n_outputs:int=3, model_name='models.resnet50'):
    """
    classification module
    :param n_classes: classes number
    :param softmax: add softmax layer
    :param feature_scale: scale value
    :param imsize: image size
    :return:
    """
    model = create_cnn(eval(model_name), nc=n_outputs, softmax=False, pretrained=True)#, y_range=[0,1])
    return model



def main():
    image_shape = (512, 256)
    a = torch.rand(2,3, image_shape[0], image_shape[1]).cuda()
    model = regressor_cnn(n_inputs=3, n_outputs=3)
    model = model.cuda()
    model = torch.nn.DataParallel(model, device_ids=range(2))

    b = model.forward(a)

    print(b.shape)


if __name__ == '__main__':
    main()